import datetime
import json
import sys
import time
import uuid
from typing import *

from aiohttp import ClientWebSocketResponse
from loguru import logger
from cryptography.hazmat.primitives.serialization import load_pem_private_key
import aiohttp
import asyncio
from threading import Lock, Thread, Event
from collections import deque, defaultdict

import config
from utils.error import WesError
from utils.tools import SingleTon


class WesExchange:

    def __init__(self, api_key: str = None, api_secret: str = None, spot: bool = True):
        self.url = 'wss://stream.binance.com/ws' if spot else 'wss://fstream.binance.com/ws'
        self._api_key: Optional[str] = api_key
        self._api_secret: Optional[str] = api_secret
        self.is_login: bool = api_key is not None and api_secret is not None

        self.session: aiohttp.ClientSession = aiohttp.ClientSession()
        self.ws: Optional[ClientWebSocketResponse] = None

        # wes任务id和回调函数
        self.callback_map: Dict[str, Callable[[Dict[str, Any]], None]] = {}
        # 已订阅事件
        self._event: Set[str] = set()
        # 订阅事件缓存
        self.event_cache: Dict[str, deque] = defaultdict(lambda: deque(maxlen=20))

        # 缓存读写锁
        self.write_lock = Lock()
        self.read_lock = Lock()
        # 准备完成信号量
        self.semaphore: Event = Event()

        self._loop = asyncio.get_event_loop()
        self.task: Optional[Thread] = None
        self.task = Thread(target=self._run, args=(self, self._loop))
        self.task.daemon = True
        self.task.start()

        self.semaphore.wait(5)
        self._running = True

    async def _send(self, method: str, data: Any, callback: Callable[[Dict[str, Any]], None] = None):
        """
        发送 websocket 报文
        :param method: 访问的binance接口
        :param data: 报文参数
        :param callback: 回调函数，按ID绑定，调用一次后删除
        :return:
        """
        _id = uuid.uuid4().__str__()
        data = {
            "method": method,
            "params": data,
            "id": _id
        }
        if self.ws:
            await self.ws.send_str(json.dumps(data))
        else:
            raise WesError('wes连接未开启')

        if callback is not None:
            self.callback_map[_id] = callback

    @staticmethod
    def _run(wex: "WesExchange", loop: asyncio.AbstractEventLoop):
        """
        监听任务后台执行
        :return:
        """
        loop.run_until_complete(wex._listen())

    async def _listen(self):
        """

        :return:
        """
        self.ws = await self.session.ws_connect(self.url, proxy=config.PROXY_CONF.get('https'))
        self.semaphore.set()
        logger.info(f'wes建立完成')

        async for msg in self.ws:
            if msg.__getattribute__('type') == aiohttp.WSMsgType.TEXT:
                data: Dict[str, Any] = json.loads(msg.__getattribute__('data'))

                # 任务回调
                if (_id := data.get('id')) is not None:
                    func = self.callback_map.pop(_id)
                    if not callable(func):
                        logger.error(f'未找到回调函数：{_id}')
                    func(data)
                # 订阅事件
                elif (event := data.get('e')) is not None:
                    self.write_lock.acquire()
                    self.event_cache[event].append(data)
                    self.write_lock.release()
                else:
                    logger.warning(f'未知响应：{data}')

            elif msg.__getattribute__('type') == aiohttp.WSMsgType.ERROR:
                logger.error(f'wes错误message: {msg}')

            else:
                logger.warning(f'wes接收到错误message：{msg.__getattribute__("type")}')

    def subscribe(self, type_: Union[str, List[str]]):
        """
        订阅事件
        :param type_:
        :return:
        """
        data = [type_] if isinstance(type_, str) else [type_]
        asyncio.run(self._send('SUBSCRIBE', data, callback=lambda x: logger.info(x)))
        logger.info(f'订阅事件：{type_}')
        self._event |= set(data)

    async def _close(self):
        await self.ws.close()
        await self.session.close()

    def shutdown(self):
        """

        :return:
        """
        if self.ws:
            self._loop.call_soon_threadsafe()
        self.task.join(timeout=10)
        if self.task.is_alive():
            logger.warning(f'wes关闭超时')
        self._running = False

    def __del__(self):
        if self._running:
            self.shutdown()
        self._loop.close()


def main():
    wes = WesExchange('12', '12')
    wes.subscribe('btcusdt@kline_1m')
    time.sleep(10)
    print(wes.event_cache['kline'][0])

    wes.shutdown()
    time.sleep(10)


if __name__ == '__main__':
    main()
